#include <ntddk.h>

#define DEVICE_NAME        L"\\Device\\NTMODELDRV_IOCONTROL"
#define SYMBOLIC_LINK_NAME L"\\??\\NTMODELDRV_IOCONTROL"

#define DRV_PREFIX "NTMODELDRV_IOCONTROL:"

#define IOCTL_NTMODELDRV_READ  CTL_CODE(0x8000, 0x800, METHOD_BUFFERED, FILE_ANY_ACCESS)
#define IOCTL_NTMODELDRV_WRITE CTL_CODE(0x8000, 0x801, METHOD_NEITHER, FILE_ANY_ACCESS)

NTSTATUS
CompleteIrp(PIRP Irp, NTSTATUS Status, ULONG_PTR Information);

NTSTATUS
NTModelDrvIoControlCreateClose(PDEVICE_OBJECT DeviceObject, PIRP Irp);

NTSTATUS
NTModelDrvIoControlDevice(PDEVICE_OBJECT DeviceObject, PIRP Irp);

VOID
NTModelDrvIoControlUnload(PDRIVER_OBJECT DriverObject);

NTSTATUS
DriverEntry(PDRIVER_OBJECT DriverObject, PUNICODE_STRING RegisterPath)
{
	NTSTATUS Status            = STATUS_SUCCESS;
	DriverObject->DriverUnload = NTModelDrvIoControlUnload;

	UNICODE_STRING usDeviceName       = { 0 };
	UNICODE_STRING usSymbolicLinkName = { 0 };
	PDEVICE_OBJECT DeviceObject       = NULL;

	do {

		RtlInitUnicodeString(&usDeviceName, DEVICE_NAME);
		RtlInitUnicodeString(&usSymbolicLinkName, SYMBOLIC_LINK_NAME);

		Status = IoCreateDevice(DriverObject, 0, &usDeviceName, FILE_DEVICE_UNKNOWN, 0, TRUE, &DeviceObject);
		if (NT_SUCCESS(Status) == FALSE) {

			KdPrintEx((77, 0, DRV_PREFIX" IoCreateDevice Failed: %x\n", Status));
			break;
		}

		Status = IoCreateSymbolicLink(&usSymbolicLinkName, &usDeviceName);
		if (NT_SUCCESS(Status) == FALSE) {

			KdPrintEx((77, 0, DRV_PREFIX" IoCreateSymbolicLink Failed: %x\n", Status));
			break;
		}

		DriverObject->MajorFunction[IRP_MJ_CREATE]         = NTModelDrvIoControlCreateClose;
		DriverObject->MajorFunction[IRP_MJ_CLOSE]          = NTModelDrvIoControlCreateClose;
		DriverObject->MajorFunction[IRP_MJ_DEVICE_CONTROL] = NTModelDrvIoControlDevice;


		KdPrintEx((77, 0, DRV_PREFIX" NTModelDrv loaded successfully\n"));

	} while (FALSE);

	if (DeviceObject != NULL) {

		if (NT_SUCCESS(Status) == FALSE) {

			IoDeleteDevice(DeviceObject);
		}
	}

	return Status;
}

VOID
NTModelDrvIoControlUnload(PDRIVER_OBJECT DriverObject)
{
	UNICODE_STRING usSymbolicLinkName = { 0 };
	RtlInitUnicodeString(&usSymbolicLinkName, SYMBOLIC_LINK_NAME);

	IoDeleteSymbolicLink(&usSymbolicLinkName);
	IoDeleteDevice(DriverObject->DeviceObject); 

	KdPrintEx((77, 0, DRV_PREFIX" NTModelDrv uninstalled successfully\n"));
}

NTSTATUS
CompleteIrp(PIRP Irp, NTSTATUS Status, ULONG_PTR Information)
{
	Irp->IoStatus.Status = Status;
	Irp->IoStatus.Information = Information;

	IoCompleteRequest(Irp, IO_NO_INCREMENT);

	return Status;
}

NTSTATUS
NTModelDrvIoControlCreateClose(PDEVICE_OBJECT DeviceObject, PIRP Irp)
{
	return CompleteIrp(Irp, STATUS_SUCCESS, 0);
}

NTSTATUS
NTModelDrvIoControlDevice(PDEVICE_OBJECT DeviceObject, PIRP Irp)
{
	PIO_STACK_LOCATION Stack = IoGetCurrentIrpStackLocation(Irp);
	
	ULONG_PTR Information = 0;
	NTSTATUS Status       = STATUS_SUCCESS;

	switch (Stack->Parameters.DeviceIoControl.IoControlCode) {

	case IOCTL_NTMODELDRV_READ: {

		ULONG_PTR Len = Stack->Parameters.DeviceIoControl.InputBufferLength;

		WCHAR szText[] = L"Hi, R3, I come for R0.";
		
		ULONG_PTR ReturnLen = Len < sizeof(szText) ? Len : sizeof(szText);

		RtlCopyMemory(Irp->AssociatedIrp.SystemBuffer, szText, ReturnLen);

		Information = ReturnLen;
		Status      = STATUS_SUCCESS;

		break;
	}

	case IOCTL_NTMODELDRV_WRITE: {

		UNICODE_STRING usText = { 0 };
		RtlInitUnicodeString(&usText, (WCHAR*)Stack->Parameters.DeviceIoControl.Type3InputBuffer);

		KdPrintEx((77, 0, DRV_PREFIX" %wZ\n", usText));

		WCHAR szText[] = L"Hi, R3, I have received your data and processed it successfully.";

		RtlCopyMemory(Irp->UserBuffer, szText, sizeof(szText));

		break;
	}

	}

	return CompleteIrp(Irp, Status, Information);
}